package com.ctfo.ulcp.exp.core;

import com.alibaba.dashscope.utils.Constants;
import dev.langchain4j.chain.ConversationalRetrievalChain;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.dashscope.QwenChatModel;
import dev.langchain4j.model.embedding.BgeSmallZhQuantizedEmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import static dev.langchain4j.data.document.loader.FileSystemDocumentLoader.loadDocument;

public class Hello {

    public static void main(String[] args) {
        String level =  System.getenv()
                .getOrDefault(Constants.DASHSCOPE_SDK_LOGGING_LEVEL_ENV, "WARN");

//        System.getenv().set

        System.out.println(level);
        if(true){
//            return;
        }
        ChatLanguageModel chatLanguageModel = QwenChatModel.builder()
                .apiKey("sk-0904c8b6cc9b4d14b71c70120b83a0fe")
                .modelName("qwen-plus")
//                .logRequests(true)
//                .logResponses(true)
                .build();

        EmbeddingModel embeddingModel = new BgeSmallZhQuantizedEmbeddingModel();

        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(DocumentSplitters.recursive(300, 0))
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .build();

        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);
        String docName = "C:\\Users\\10003021\\Documents\\qf2022.pdf";
        Document document = loadDocument(docName, new ApachePdfBoxDocumentParser());
        ingestor.ingest(document);





        ConversationalRetrievalChain chain = ConversationalRetrievalChain.builder()
                .chatLanguageModel(chatLanguageModel)
                .chatMemory(chatMemory)
                .contentRetriever(new EmbeddingStoreContentRetriever(embeddingStore, embeddingModel,10))
//                    .retriever(EmbeddingStoreRetriever.from(embeddingStore, embeddingModel))
                // .chatMemory() // you can override default chat memory
                // .promptTemplate() // you can override default prompt template
                .build();

        String[] questsons = new String[]{
                "千方科技营收是多少？",
                "千方科技主营业务是什么？",
                "千方科技实控人是谁？",
                "总部在哪里？",
                "营收同比增长多少？",
        };
        for (String qustion: questsons) {
            String answer = chain.execute(qustion);
            System.out.println(answer);
        }
    }

}
